import cPickle
import string
from random import Random

class UndoOperationDoesNotExist(KeyError):
    "Used to indicate that the requested undo operation with the specific id does not exist"

def save_undo(request, model_instance, redirect_to=None):
    """
    Saves an undo operation.
    Returns the id of the undo operation.
    If the optional redirect_to argument is not supplied, the model_instance must implement
    a get_absolute_url() method.

    This function saves the model_instance and all its related objects.
    """

    # create a list of the the model instance and all objects that relate to the model instance
    seen_objs = {}
    model_instance._collect_sub_objects(seen_objs)

    ordered_classes = seen_objs.keys()
    ordered_classes.reverse()

    objs = []
    for cls in ordered_classes:
        seen_objs[cls] = seen_objs[cls].items()

        for pk, instance in seen_objs[cls]:
            objs.append(instance)

    # serialise the list of model instances
    pickle = cPickle.dumps(objs)

    # generate a random string for use as undo id
    id = u''.join( Random().sample(string.letters+string.digits, 12) )

    # save the undo operation in the session
    request.session['undo'] = {'pickle':pickle,
                               'id':id,
                               'redirect_to':redirect_to if redirect_to else model_instance.get_absolute_url()}

    return id


def do_undo(request, undo_id):
    """
    Executes an undo operation.
    Returns the url to redirect to, after the undo operation has executed.
    Raises UndoOperationDoesNotExist if undo_id does not match the currently saved undo operation.

    It is recommended to use the supplied generic undo view in django_undo.views instead of this function.
    """

    try:
        undo = request.session['undo']
        del request.session['undo']

        # check if the undo operations id matches the requested id
        if undo['id'] != undo_id:
            raise KeyError

    except KeyError:
        raise UndoOperationDoesNotExist, u'The undo operation does not exist'

    objects = cPickle.loads(undo['pickle'])
    for obj in objects:
        obj.save()

    return undo['redirect_to']
